// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#include <stdarg.h>
#include "ppl/common/mmap.h"
#include "ppl/nn/common/logger.h"
#include "ppl/nn/optimizers/utils.h"
#include "ppl/nn/optimizers/engine_graph_partitioner.h"
#include "ppl/nn/runtime/runtime_impl.h"
#include "ppl/nn/models/onnx/model_parser.h"
#include "ppl/nn/models/onnx/runtime_builder_impl.h"
using namespace std;
using namespace ppl::common;

#ifdef PPLNN_ENABLE_PMX_MODEL
#include "ppl/nn/models/pmx/serializer.h"
#endif

namespace ppl { namespace nn { namespace onnx {

RuntimeBuilderImpl::RuntimeBuilderImpl() {
    graph_info_ = make_shared<RuntimeGraphInfo>();
    aux_info_ = make_shared<RuntimeAuxInfo>();
}

RuntimeBuilderImpl::~RuntimeBuilderImpl() {
    aux_info_.reset();
    graph_info_.reset();
}

RetCode RuntimeBuilderImpl::LoadModel(const char* model_buf, uint64_t buf_len, const char* model_file_dir) {
    auto status = ModelParser::Parse(model_buf, buf_len, model_file_dir, &model_);
    if (status != RC_SUCCESS) {
        LOG(ERROR) << "parse graph failed: " << GetRetCodeStr(status);
        return status;
    }

    return RC_SUCCESS;
}

RetCode RuntimeBuilderImpl::LoadModel(const char* model_file) {
    Mmap fm;
    auto status = fm.Init(model_file, Mmap::READ);
    if (status != RC_SUCCESS) {
        LOG(ERROR) << "mapping file [" << model_file << "] faild.";
        return status;
    }

    string parent_dir;
    auto pos = string(model_file).find_last_of("/\\");
    if (pos == string::npos) {
        parent_dir = ".";
    } else {
        parent_dir.assign(model_file, pos);
    }

    return LoadModel(fm.GetData(), fm.GetSize(), parent_dir.c_str());
}

RetCode RuntimeBuilderImpl::LoadModel(const char* model_buf, uint64_t buf_len, const char** inputs, uint32_t nr_input,
                                      const char** outputs, uint32_t nr_output, const char* model_file_dir) {
    auto status = ModelParser::Parse(model_buf, buf_len, model_file_dir, inputs, nr_input, outputs, nr_output, &model_);
    if (status != RC_SUCCESS) {
        LOG(ERROR) << "parse graph failed: " << GetRetCodeStr(status);
        return status;
    }

    return RC_SUCCESS;
}

RetCode RuntimeBuilderImpl::LoadModel(const char* model_file, const char** inputs, uint32_t nr_input,
                                      const char** outputs, uint32_t nr_output) {
    Mmap fm;
    auto status = fm.Init(model_file, Mmap::READ);
    if (status != RC_SUCCESS) {
        LOG(ERROR) << "mapping file [" << model_file << "] faild.";
        return status;
    }

    string parent_dir;
    auto pos = string(model_file).find_last_of("/\\");
    if (pos == string::npos) {
        parent_dir = ".";
    } else {
        parent_dir.assign(model_file, pos);
    }

    return LoadModel(fm.GetData(), fm.GetSize(), inputs, nr_input, outputs, nr_output, parent_dir.c_str());
}

RetCode RuntimeBuilderImpl::SetResources(const Resources& resource) {
    resource_.engines.resize(resource.engine_num);
    for (uint32_t i = 0; i < resource.engine_num; ++i) {
        resource_.engines[i] = static_cast<EngineImpl*>(resource.engines[i]);
    }

    resource_.graph_partitioner = make_shared<EngineGraphPartitioner>();
    return RC_SUCCESS;
}

RetCode RuntimeBuilderImpl::Preprocess() {
    auto status = utils::ProcessGraph(resource_, &model_.graph, graph_info_.get());
    if (status != RC_SUCCESS) {
        LOG(ERROR) << "process graph failed: " << GetRetCodeStr(status);
        return status;
    }

    status = aux_info_->Init(model_.graph.topo.get(), resource_.reserved_edgeids);
    if (status != RC_SUCCESS) {
        LOG(ERROR) << "GenerateRuntimeAuxInfo failed: " << GetRetCodeStr(status);
        return status;
    }

    return RC_SUCCESS;
}

Runtime* RuntimeBuilderImpl::CreateRuntime() const {
    auto runtime = new RuntimeImpl();
    if (!runtime) {
        return nullptr;
    }

    auto status = runtime->Init(model_.graph.topo, graph_info_, aux_info_, resource_.reserved_edgeids);
    if (status != RC_SUCCESS) {
        LOG(ERROR) << "init runtime failed: " << GetRetCodeStr(status);
        delete runtime;
        return nullptr;
    }

    return runtime;
}

RetCode RuntimeBuilderImpl::Serialize(const char* fmt, const void* options, utils::DataStream* ds) const {
#ifdef PPLNN_ENABLE_PMX_MODEL
    if (fmt == string("pmx")) {
        const pmx::SaveModelOptions default_opt;
        const pmx::SaveModelOptions* opt;
        if (options) {
            opt = (const pmx::SaveModelOptions*)options;
        } else {
            opt = &default_opt;
        }
        pmx::Serializer serializer;
        return serializer.Serialize(*opt, model_.graph.topo.get(), resource_.engines, *graph_info_, ds);
    }
#endif
    LOG(ERROR) << "model format[" << fmt << "] is not supported.";
    return RC_UNSUPPORTED;
}

RetCode RuntimeBuilderImpl::ReserveTensor(const char* tensor_name) {
    auto edge = model_.graph.topo->GetEdge(tensor_name);
    if (!edge) {
        LOG(ERROR) << "ReserveTensor: cannot find tensor named[" << tensor_name << "]";
        return RC_NOT_FOUND;
    }

    resource_.reserved_edgeids.insert(edge->GetId());
    return RC_SUCCESS;
}

}}} // namespace ppl::nn::onnx
